#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import numpy as np
import gi
import time

gi.require_version('Gst', '1.0')
from gi.repository import Gst

class GstVideoPublisher(Node):
    def __init__(self):
        super().__init__('gst_video_publisher')
        
        # 参数配置
        self.declare_parameters(
            namespace='',
            parameters=[
                ('udp_address', '230.1.1.1'),
                ('udp_port', 1720),
                ('interface', 'enx00e04c69101b'),
                ('target_fps', 30),
                ('image_topic', '/camera/image_raw'),
                ('width', 1280),
                ('height', 720),
                ('frame_id', 'camera_optical_frame')
            ]
        )
        
        # 获取参数值
        self.udp_address = self.get_parameter('udp_address').value
        self.udp_port = self.get_parameter('udp_port').value
        self.interface = self.get_parameter('interface').value
        self.target_fps = self.get_parameter('target_fps').value
        self.image_topic = self.get_parameter('image_topic').value
        self.width = self.get_parameter('width').value
        self.height = self.get_parameter('height').value
        self.frame_id = self.get_parameter('frame_id').value
        
        # 初始化ROS2 Publisher
        self.publisher = self.create_publisher(Image, self.image_topic, 10)
        self.bridge = CvBridge()
        
        # 帧率控制
        self.frame_interval = 1.0 / self.target_fps if self.target_fps > 0 else 0
        self.get_logger().info(f'目标帧率: {self.target_fps} FPS')
        self.get_logger().info(f'图像话题: {self.image_topic}')
        self.get_logger().info(f'视频分辨率: {self.width}x{self.height}')
        
        # GStreamer初始化
        Gst.init(None)
        self.pipeline = None
        self.appsink = None
        self.init_gstreamer()
        
        # 帧统计
        self.frame_count = 0
        self.last_report_time = time.time()

    def init_gstreamer(self):
        """初始化GStreamer管道"""
        pipeline_str = (
            f"udpsrc address={self.udp_address} port={self.udp_port} "
            f"multicast-iface={self.interface} ! "
            "queue ! "
            "application/x-rtp,media=video,encoding-name=H264 ! "
            "rtph264depay ! h264parse ! avdec_h264 ! "
            "videoconvert ! "
            f"video/x-raw,width={self.width},height={self.height},format=BGR,framerate={self.target_fps}/1 ! "
            "appsink name=sink emit-signals=True drop=True sync=false"
        )
        
        self.get_logger().info(f'GStreamer管道配置:\n{pipeline_str}')
        
        try:
            self.pipeline = Gst.parse_launch(pipeline_str)
            self.appsink = self.pipeline.get_by_name("sink")
            self.appsink.connect("new-sample", self.on_new_frame, None)
            self.pipeline.set_state(Gst.State.PLAYING)
            self.get_logger().info('GStreamer管道启动成功')
        except Exception as e:
            self.get_logger().error(f'GStreamer初始化失败: {str(e)}')
            raise

    def on_new_frame(self, sink, _):
        """GStreamer新帧回调函数"""
        sample = sink.emit("pull-sample")
        if not sample:
            self.get_logger().warning("收到空样本", throttle_duration_sec=5)
            return Gst.FlowReturn.OK

        buf = sample.get_buffer()
        if not buf:
            self.get_logger().warning("空缓冲区", throttle_duration_sec=5)
            return Gst.FlowReturn.OK

        success, map_info = buf.map(Gst.MapFlags.READ)
        if not success:
            self.get_logger().error("缓冲区映射失败")
            return Gst.FlowReturn.ERROR

        try:
            # 方法一：直接构造numpy数组（更高效）
            img_array = np.ndarray(
                shape=(self.height, self.width, 3),
                dtype=np.uint8,
                buffer=map_info.data
            )
            
            # 方法二（备用）：如果方法一无效，使用reshape
            # img_array = np.frombuffer(map_info.data, dtype=np.uint8).reshape(
            #     (self.height, self.width, 3)
            # )
            
            # 调试：检查图像数据
            if img_array.shape != (self.height, self.width, 3):
                self.get_logger().error(
                    f"图像尺寸异常: 期望 {self.height}x{self.width}x3, 实际 {img_array.shape}",
                    throttle_duration_sec=1
                )
                return Gst.FlowReturn.ERROR
                
            # 发布ROS消息
            ros_img = self.bridge.cv2_to_imgmsg(img_array, "bgr8")
            ros_img.header.stamp = self.get_clock().now().to_msg()
            ros_img.header.frame_id = self.frame_id
            self.publisher.publish(ros_img)
            
            # 帧率统计
            self.frame_count += 1
            now = time.time()
            if now - self.last_report_time >= 5.0:
                actual_fps = self.frame_count / (now - self.last_report_time)
                self.get_logger().info(
                    f'实际帧率: {actual_fps:.2f} FPS | '
                    f'分辨率: {img_array.shape[1]}x{img_array.shape[0]}'
                )
                self.frame_count = 0
                self.last_report_time = now

        except Exception as e:
            self.get_logger().error(f'帧处理异常: {str(e)}')
        finally:
            buf.unmap(map_info)
        
        return Gst.FlowReturn.OK

    def destroy(self):
        """资源清理"""
        if self.pipeline:
            self.pipeline.set_state(Gst.State.NULL)
            self.get_logger().info("GStreamer管道已停止")
        super().destroy_node()

def main(args=None):
    rclpy.init(args=args)
    node = GstVideoPublisher()
    
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        node.get_logger().info("节点被手动终止")
    except Exception as e:
        node.get_logger().error(f'节点运行时错误: {str(e)}')
    finally:
        node.destroy()
        rclpy.shutdown()

if __name__ == '__main__':
    main()